import asyncio
import websockets
import json
import base64
import cv2
import numpy as np
import time

class PoseDetectionClient:
    # def __init__(self, uri='ws://192.168.58.125:8889'):
    def __init__(self, uri='ws://localhost:8889'):
    # 如果需要修改服务器地址，请更改上面的uri
        self.uri = uri
        self.websocket = None
        self.poses = None
        self.center_poses = None
        self.connected = False
        self.ready = False
        
    async def connect(self, max_retries=5, retry_delay=2):
        """连接到服务器，支持重试"""
        for attempt in range(max_retries):
            try:
                print(f"尝试连接 {self.uri} (第 {attempt + 1} 次)")
                self.websocket = await websockets.connect(
                    self.uri, 
                    ping_timeout=120,  # 增加ping超时到120秒
                    ping_interval=20, # 每20秒发送一次ping
                    close_timeout=30  # 关闭超时30秒
                )
                self.connected = True
                print(f"成功连接到 {self.uri}")
                return True
            except ConnectionRefusedError as e:
                print(f"连接被拒绝 (尝试 {attempt + 1}/{max_retries}): {e}")
                if attempt < max_retries - 1:
                    print(f"等待 {retry_delay} 秒后重试...")
                    await asyncio.sleep(retry_delay)
                else:
                    print("达到最大重试次数，连接失败")
                    raise e
            except Exception as e:
                print(f"连接错误 (尝试 {attempt + 1}/{max_retries}): {e}")
                if attempt < max_retries - 1:
                    print(f"等待 {retry_delay} 秒后重试...")
                    await asyncio.sleep(retry_delay)
                else:
                    print("达到最大重试次数，连接失败")
                    raise e
        return False
    
    async def detect_poses(self, is_first_frame=False):
        """请求位姿检测"""
        message = {
            'command': 'detect',
            'is_first_frame': is_first_frame,
            'timestamp': time.time()
        }
        
        await self.websocket.send(json.dumps(message))
        response = await self.websocket.recv()
        data = json.loads(response)
        
        if data['status'] == 'success':
            # 解码原始图像
            orig_image_data = base64.b64decode(data['original_image'])
            orig_nparr = np.frombuffer(orig_image_data, np.uint8)
            original_image = cv2.imdecode(orig_nparr, cv2.IMREAD_COLOR)
            
            # 解码可视化图像
            vis_image_data = base64.b64decode(data['visualization_image'])
            vis_nparr = np.frombuffer(vis_image_data, np.uint8)
            visualization_image = cv2.imdecode(vis_nparr, cv2.IMREAD_COLOR)
            
            return {
                'poses': np.array(data['poses']),
                'center_poses': np.array(data['center_poses']),
                'original_image': original_image,
                'visualization_image': visualization_image,
                'timestamp': data['timestamp']
            }
        else:
            raise Exception(data['message'])
    
    async def close(self):
        """关闭连接"""
        if self.websocket:
            await self.websocket.close()
            self.connected = False
            print("连接已关闭")
    
    def is_connected(self):
        """检查连接状态"""
        return self.connected and self.websocket is not None

async def main():
    client = PoseDetectionClient()
    
    try:
        # 尝试连接服务器
        success = await client.connect()
        if not success:
            print("无法连接到服务器，请检查：")
            print("1. 服务器是否运行在 192.168.58.125:8889")
            print("2. 防火墙是否允许8889端口")
            print("3. 网络连接是否正常")
            return

        print("开始位姿检测循环...")
        while True:
            try:
                # 检测位姿
                result = await client.detect_poses(is_first_frame=False)
                
                # 处理结果
                print(f"Poses: {result['poses'].shape}")
                print(f"Center poses: {result['center_poses'].shape}")
                
                # 显示图像
                cv2.imshow('Remote Pose Detection', result['visualization_image'])
                
                # 按键处理
                key = cv2.waitKey(1) & 0xFF
                if key == ord(' '):  # 空格键重新初始化
                    print("redetect...")
                    result = await client.detect_poses(is_first_frame=True)
                elif key == ord('c'):
                    client.ready = True
                elif key == 27:  # ESC键退出
                    break
                    
                # 控制帧率
                await asyncio.sleep(0.033)  # ~30 FPS
                
            except Exception as e:
                print(f"Detection error: {e}")
                break
                    
    finally:
        print("Closing connection...")
        await client.close()
        cv2.destroyAllWindows()

if __name__ == "__main__":
    asyncio.run(main())